import sys
import glob
import os

FindStri = sys.argv[1]  #like:P1,P2
FindFlag = sys.argv[2]  #like:CDNA,TCR,TCRDG,BCR
workdir = sys.argv[3]
outdir = sys.argv[4]

#python Auto-FindSCSamples.py P1,P2 TCRDG /data2/users/dongxingjun/Human_SC_AnaData

#finish in 2020.9.27
#output: Px_Tumor_	/data/xxx

#test
#FindStri = "P1,P2"
#FindFlag = "TCR"
#workdir = "/data1/home/zhoulab/HumanSC-Data"
#outdir = "/data1/home/zhoulab"

#function
def getPathFile(basedir,sample,source,atype):
    apath = False
    if atype == "CDNA":
        apath = "%s/%s/GRCh38" % (basedir,atype)
        if os.path.isdir(apath):
            stri = "%s_%s_\t%s\tF" % (sample,source,apath)
            return stri

    elif atype == "TCR":
        apath = "%s/%s/outs" % (basedir,atype)
        if os.path.isdir(apath):
            stri = "%s_%s_\t%s" % (sample,source,apath)
            return stri

    elif atype == "BCR":
        apath = "%s/%s/outs" % (basedir,atype)
        if os.path.isdir(apath):
            stri = "%s_%s_\t%s" % (sample,source,apath)
            return stri

    print("error:",basedir,sample,atype)
    return ""

    
#function

#main

#
paths = "%s/P*" % (workdir)

#pat dir
AllPatsPath = glob.glob(paths)
AllPatsPathDict = {}
AllPatsList = []
for i in AllPatsPath:
    if os.path.isdir(i):
        Pat = os.path.basename(i)
        AllPatsPathDict[Pat] = i
        AllPatsList.append(Pat)
#pat dir

#every pat type
AllPatsTypeDict = {}
for i in AllPatsList:
    ThisPath = "%s/*" % (AllPatsPathDict[i])
    ThisTypesPath = glob.glob(ThisPath)

    PatTypes = []
    for y in ThisTypesPath:
        ThisT = os.path.basename(y)
        PatTypes.append(ThisT)

    AllPatsTypeDict[i] = PatTypes
#every pat type

if FindStri != "All":   #P1,P2,P3
    FindStriTemp = FindStri.split(",")

    for i in FindStriTemp: #every patient
        outFile = "%s/%s.%s.list" % (outdir,i,FindFlag)
        w = open(outFile,mode = "w+")
        for y in AllPatsTypeDict[i]: #Tumor,Blood,Normal,...
            adir = "%s/%s" % (AllPatsPathDict[i],y)
            thisline = getPathFile(adir,i,y,FindFlag)
            if thisline != "":
                thisline += "\n"
                w.write(thisline)
        w.close()

